package fem2.observer;

import inf.jlinalg.ArrayVector;
import inf.jlinalg.IVector;

import java.awt.GridLayout;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;

import javax.swing.JFrame;
import javax.swing.SwingUtilities;

import ptolemy.plot.Plot;
import fem2.Element;
import fem2.Model;
import fem2.Node;
import fem2.Observer;
import fem2.strategies.NewtonRaphsonStrategy;
import fem2.strategies.Strategy;

/**
 * A resistant plotter for structural application. If multiple nodes are
 * provided, it will sum all the resistant force at nodes
 * 
 * @author hbui
 * 
 */
public class ResistantGraphicalObserver extends JFrame implements Observer {

	private static final long serialVersionUID = 1L;

	/**
	 * class represent a pair of element and an associated id with this element
	 * 
	 * @author hbui
	 * 
	 */
	private class Pair {
		private Element e;
		private int id;

		Pair(Element e, int id) {
			this.e = e;
			this.id = id;
		}

		public Element getElement() {
			return this.e;
		}

		public int getId() {
			return this.id;
		}
	}

	private NewtonRaphsonStrategy s;

	private Map<Node, ArrayList<Pair>> mapNodeToElem = new HashMap<Node, ArrayList<Pair>>();

	private Plot plotTime = new Plot();
	private double step;
	private int dim;
	private boolean[] en;
	private boolean inverse = false;
	private Writer fileWriter = null;
	private boolean setWrite = false;

	/**
	 * 
	 * @param m
	 * @param s
	 * @param nodes
	 */
	public ResistantGraphicalObserver(Model m, Strategy s, Node... nodes) {
		this(m, s, 7, nodes);
	}

	/**
	 * 
	 * @param m
	 * @param s
	 * @param mode
	 *            mode = 1 : plot x resistant<br>
	 *            mode = 2 : plot y resistant<br>
	 *            mode = 3 : plot z resistant<br>
	 *            mode = 4 : plot xy resistant<br>
	 *            mode = 5 : plot yz resistant<br>
	 *            mode = 6 : plot xz resistant<br
	 *            mode = 7 : plot xyz resistant
	 * @param nodes
	 */
	public ResistantGraphicalObserver(Model m, Strategy s, int mode, Node... nodes) {
		if (s != null) {
			if (!(s instanceof NewtonRaphsonStrategy)) {
				throw new Error("this observer only support NewtonRaphsonBasedStrategy");
			} else {
				this.s = (NewtonRaphsonStrategy) s;
			}
		}
		/*
		 * build map of element (with associated internal index) for nodes
		 */
		if (nodes != null) {
			dim = m.getMesh().getDimension();
			for (Node n : nodes) {
				buidMap(m, n);
			}
		}
		/*
		 * enable modes
		 */
		en = new boolean[dim];
		for (int i = 0; i < dim; i++) {
			en[i] = false;
		}
		if (mode == 1 || mode == 4 || mode == 6 || mode == 7) {
			en[0] = true;
		}
		if (mode == 2 || mode == 4 || mode == 5 || mode == 7) {
			en[1] = true;
		}
		if (mode == 3 || mode == 5 || mode == 6 || mode == 7 && dim == 3) {
			en[2] = true;
		}
		/*
		 * set up step
		 */
		step = -1.0;
		/*
		 * set up the plot
		 */
		this.setLayout(new GridLayout());
		this.add(this.plotTime);
		this.plotTime.setTitle("Resistant plot");
		this.pack();
	}

	/**
	 * set the inversion of the graph
	 * 
	 * @param inverse
	 */
	public void setInverse(boolean inverse) {
		this.inverse = inverse;
	}

	/**
	 * set the output to file
	 * 
	 * @param fileName
	 */
	public void setOutputToFile(String fileName) throws IOException {
		File file = new File(fileName);
		if (!file.exists()) {
			/*
			 * delete existing file
			 */
			file.delete();
			fileWriter = new BufferedWriter(new FileWriter(file, true));
		}
		/*
		 * create new file
		 */
		fileWriter = new BufferedWriter(new FileWriter(file));
		setWrite = true;
	}

	/**
	 * add the origin to the resistant plot
	 */
	public void addOrigin() {
		for (int i = 0; i < dim; i++) {
			if (en[i]) {
				this.plotTime.addPoint(i, 0, 0, true);
			}
		}
	}

	@Override
	public void fireStarted() {
		this.setVisible(true);
	}

	@Override
	public void fireTimeStepStarted() {
	}

	@Override
	public void fireIterationStarted() {
	}

	@Override
	public void fireIterationFinished() {
	}

	@Override
	public void fireTimeStepFinished() {
		/*
		 * compute the current step
		 */
		if (s != null) {
			step = s.getCurrentTimeStep();
		} else {
			step += 1.0;
		}
		if (setWrite) {
			try {
				fileWriter.write(Double.toString(step) + "\n");
			} catch (IOException e) {
				e.printStackTrace();
				throw new Error("Write failed");
			}
		}
		/*
		 * processing data
		 */
		double[] re = new double[dim];
		for (Node n : mapNodeToElem.keySet()) {
			ArrayList<Pair> pairs = mapNodeToElem.get(n);
			/*
			 * compute internal force based on contribution from elements
			 */
			for (Pair p : pairs) {
				Element e = p.getElement();
				int id = p.getId();
				IVector rie = new ArrayVector(e.getMeshPart().countDofs());
				rie = (IVector) e.computeRie(rie);
				for (int i = 0; i < dim; i++) {
					re[i] += -rie.get(id * dim + i);
				}
			}
		}
		/*
		 * add to plot
		 */
		boolean connected = true;
		for (int i = 0; i < dim; i++) {
			if (en[i]) {
				if (inverse) {
					this.plotTime.addPoint(i, step, -re[i], connected);
					/*
					 * write to file of needed
					 */
					if (setWrite) {
						try {
							fileWriter.write(Double.toString(-re[i]) + "\n");
						} catch (IOException e) {
							e.printStackTrace();
							throw new Error("Write failed");
						}
					}
				} else {
					this.plotTime.addPoint(i, step, re[i], connected);
					/*
					 * write to file of needed
					 */
					if (setWrite) {
						try {
							fileWriter.write(Double.toString(re[i]) + "\n");
						} catch (IOException e) {
							e.printStackTrace();
							throw new Error("Write failed");
						}
					}
				}
			}
		}
		if (setWrite) {
			try {
				fileWriter.write("\n");
			} catch (IOException e) {
				e.printStackTrace();
				throw new Error("Write failed");
			}
		}
		/*
		 * refresh the painter
		 */
		SwingUtilities.invokeLater(new Runnable() {

			@Override
			public void run() {
				plotTime.repaint();
			}
		});
	}

	@Override
	public void fireFinished() {
		if (setWrite) {
			try {
				fileWriter.close();
			} catch (IOException e) {
				e.printStackTrace();
				throw new Error("Close failed");
			}
		}
	}

	/**
	 * build the map of element contains the node
	 * 
	 * @param m
	 * @param n
	 */
	private void buidMap(Model m, Node n) {
		ArrayList<Pair> map = new ArrayList<Pair>();
		boolean isEmpty = true;
		for (int i = 0; i < m.countElements(); i++) {
			Element e = m.getElement(i);
			int id = e.getMeshPart().contains(n);
			if (id > -1) {
				Pair p = new Pair(e, id);
				map.add(p);
				isEmpty = false;
			}
		}
		if (!isEmpty) {
			mapNodeToElem.put(n, map);
		}
	}
}
